package mr

import (
	"encoding/binary"
	"fmt"
	"hash/fnv"
	"io"
	"log"
	"net"
	"os"
	"strconv"
	"sync"
	"time"
)

type sStatus int

const (
	Available sStatus = iota
	Busy
)

type Coordinator struct {
	// Your definitions here.
	workers sync.Map
}

// Your code here -- RPC handlers for the worker to call.

func (this *Coordinator) heartBeat(l net.Listener) {
	defer l.Close()

	for {
		// get connections
		conn, err := l.Accept()
		if err != nil {
			log.Fatalln(err)
		}

		// get heartbeat
		go func(c net.Conn) {
			defer c.Close()
			for {
				buf := make([]byte, 4)
				_, err := c.Read(buf)
				if err != nil && err != io.EOF {
					c.Close()
					break
				}
				if err == io.EOF {
					break
				}
				pid := int(binary.BigEndian.Uint32(buf))
				_, ok := this.workers.LoadOrStore(pid, Slave{id: pid, state: Available})
				if !ok {
					this.PrintWorkers()
				}
				c.Write([]byte{0})
			}
		}(conn)
	}
}

func (c *Coordinator) PrintWorkers() {
	c.workers.Range(func(k, v interface{}) bool {
		fmt.Println(k, v)
		return true
	})
}

// use ihash(key) % NReduce to choose the reduce
// task number for each KeyValue emitted by Map.
func (this *Coordinator) hash(key string) int {
	h := fnv.New32a()
	h.Write([]byte(key))
	return int(h.Sum32() & 0x7fffffff)
}

// start a thread that listens for RPCs from worker.go
func (this *Coordinator) server() {
	//l, e := net.Listen("tcp", ":1234")
	sockname := getSockName(-1)
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatalln("listen error:", e)
	}

	go this.heartBeat(l)
}

// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
func (c *Coordinator) Done() bool {
	ret := false

	// Your code here.

	return ret
}

func (c *Coordinator) pickAvailableWorkers(maxN int) []int {
	workers := make([]int, 0)
	for len(workers) == 0 {
		c.workers.Range(func(k, v interface{}) bool {
			if len(workers) == maxN {
				return false
			}
			if v.(Slave).state == Available {
				workers = append(workers, k.(int))
			}
			return true
		})
	}

	return workers
}

// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	c := Coordinator{}

	// Your code here.
	c.server()

	// wait for the workers
	time.Sleep(time.Second * 10)

	// map init
	mapers := c.pickAvailableWorkers(len(os.Args[2:]))
	argsMaps := make([]MapArgs, len(mapers))
	replyMaps := make([]MapReply, len(mapers))
	for i := range mapers {
		argsMaps[i].Filenames = make([]string, 0)
		replyMaps[i].Buckets = make(map[string]string, 0)
	}

	// map
	var wg sync.WaitGroup
	var ok bool
	wg.Add(len(mapers))
	for _, filename := range os.Args[2:] {
		key := c.hash(filename) % len(mapers)
		argsMaps[key].Filenames = append(argsMaps[key].Filenames, filename)
	}
	for i := range argsMaps {
		go func(idx int) {
			defer wg.Done()
			pid := mapers[idx]
			// _, ok := c.workers.Store(pid, Slave{id: pid, state: Busy})
			ok = rpcCall(pid, "Slave.Map", &argsMaps[idx], &replyMaps[idx])
			if ok {
				// _, ok = c.workers.Store(pid, Slave{id: pid, state: Available})
			}
		}(i)
	}
	wg.Wait()

	// reduce init
	reducers := c.pickAvailableWorkers(nReduce)
	argsReduces := make([]ReduceArgs, len(reducers))
	replyReduces := make([]ReduceReply, len(reducers))
	for i := range reducers {
		argsReduces[i].KeysFiles = make(map[string][]string, 0)
		replyReduces[i].N = make([]string, 0)
	}

	for _, i := range replyMaps {
		for k, v := range i.Buckets {
			key := c.hash(k) % len(reducers)
			argsReduces[key].KeysFiles[k] = append(argsReduces[key].KeysFiles[k], v)
		}
	}

	// reduce
	for i := range argsReduces {
		go func(idx int) {
			pid := reducers[idx]
			// _, ok := c.workers.Store(pid, Slave{id: pid, state: Busy})
			ok = rpcCall(pid, "Slave.Reduce", &argsReduces[idx], &replyReduces[idx])
			if ok {
				// _, ok = c.workers.Store(pid, Slave{id: pid, state: Available})
				oname := "mr-out-"
				oname += strconv.Itoa(idx)
				ofile, _ := os.Create(oname)
				j := 0
				for k := range argsReduces[idx].KeysFiles {
					fmt.Fprintf(ofile, "%v %v\n", k, replyReduces[idx].N[j])
					j++
				}

			}
		}(i)
	}

	return &c
}
